{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "from pymilvus import MilvusClient, DataType\n",
    "import time\n",
    "\n",
    "def create_collection(collection_name):\n",
    "    # 检查同名集合是否存在，如果存在则删除\n",
    "    if milvus_client.has_collection(collection_name):\n",
    "        print(f\"集合 {collection_name} 已经存在\")\n",
    "        try:\n",
    "            # 删除同名集合\n",
    "            milvus_client.drop_collection(collection_name)\n",
    "            print(f\"删除集合：{collection_name}\")\n",
    "        except Exception as e:\n",
    "            print(f\"删除集合时出现错误: {e}\")\n",
    "    # 创建集合模式\n",
    "    schema = MilvusClient.create_schema(\n",
    "        auto_id=True,\n",
    "        enable_dynamic_field=True,\n",
    "        # 设置分区数量，默认为16\n",
    "        num_partitions=16,\n",
    "        description=\"\"\n",
    "    )\n",
    "    # 添加字段到schema\n",
    "    schema.add_field(field_name=\"id\", datatype=DataType.INT64, is_primary=True, max_length=256)\n",
    "    schema.add_field(field_name=\"text\", datatype=DataType.VARCHAR, max_length=256)\n",
    "    # bm25稀疏向量\n",
    "    schema.add_field(field_name=\"sparse_vectors_bm25\", datatype=DataType.SPARSE_FLOAT_VECTOR)\n",
    "    # splade稀疏向量\n",
    "    schema.add_field(field_name=\"sparse_vectors_splade\", datatype=DataType.SPARSE_FLOAT_VECTOR)\n",
    "    # 创建集合\n",
    "    try:\n",
    "        milvus_client.create_collection(\n",
    "            collection_name=collection_name,\n",
    "            schema=schema,\n",
    "            shards_num=2\n",
    "        )\n",
    "        print(f\"创建集合：{collection_name}\")\n",
    "    except Exception as e:\n",
    "        print(f\"创建集合的过程中出现了错误: {e}\")\n",
    "    # 等待集合创建成功\n",
    "    while not milvus_client.has_collection(collection_name):\n",
    "        # 获取集合的详细信息\n",
    "        time.sleep(1)\n",
    "    if milvus_client.has_collection(collection_name):\n",
    "        print(f\"集合 {collection_name} 创建成功\")\n",
    "\n",
    "# 示例\n",
    "collection_name = \"docs\"\n",
    "uri=\"http://localhost:19530\"\n",
    "milvus_client = MilvusClient(uri=uri)\n",
    "create_collection(collection_name)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 准备数据\n",
    "entities = [\n",
    "    {\n",
    "        # 文本字段\n",
    "        \"text\": docs[i],\n",
    "        \"text_en\": docs_en[i],\n",
    "        # bm25稀疏向量字段\n",
    "        \"sparse_vectors_bm25\": list(sparse_vectors_bm25)[i].reshape(1, -1),\n",
    "        # splade稀疏向量字段\n",
    "        \"sparse_vectors_splade\": list(sparse_vectors_splade)[i].reshape(1, -1),\n",
    "    }\n",
    "    for i in range(len(docs))\n",
    "]\n",
    "\n",
    "# 导入数据\n",
    "milvus_client.insert(collection_name=collection_name, data=entities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 创建索引参数\n",
    "index_params = milvus_client.prepare_index_params()\n",
    "\n",
    "# 为稀疏向量bm25创建索引参数\n",
    "index_params.add_index(\n",
    "    index_name=\"sparse_vectors_bm25\",\n",
    "    field_name=\"sparse_vectors_bm25\",\n",
    "    # SPARSE_INVERTED_INDEX是传统的倒排索引，SPARSE_WAND使用Weak-AND算法来减少搜索过程中的完整IP距离计算\n",
    "    index_type=\"SPARSE_INVERTED_INDEX\",\n",
    "    # 目前仅支持IP\n",
    "    metric_type=\"IP\",\n",
    "    # 创建索引时，排除向量值最小的20%的向量。对于稀疏向量来说，向量值越大，说明在该维度上的重要性越大。范围[0,1]。\n",
    "    params={\"drop_ratio_build\": 0.2}\n",
    ")\n",
    "\n",
    "\n",
    "# 为稀疏向量splade创建索引参数\n",
    "index_params.add_index(\n",
    "    index_name=\"sparse_vectors_splade\",\n",
    "    field_name=\"sparse_vectors_splade\",\n",
    "    # SPARSE_INVERTED_INDEX是传统的倒排索引，SPARSE_WAND使用Weak-AND算法来减少搜索过程中的完整IP距离计算\n",
    "    index_type=\"SPARSE_INVERTED_INDEX\",\n",
    "    # 目前仅支持IP\n",
    "    metric_type=\"IP\",\n",
    "    # 创建索引时，排除向量值最小的20%的向量。对于稀疏向量来说，向量值越大，说明在该维度上的重要性越大。范围[0,1]。\n",
    "    params={\"drop_ratio_build\": 0.2}\n",
    ")\n",
    "\n",
    "# 创建索引\n",
    "milvus_client.create_index(\n",
    "    collection_name=collection_name,\n",
    "    index_params=index_params\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 查看索引信息\n",
    "def show_index_info(collection_name: str) -> None:\n",
    "    \"\"\"\n",
    "    显示指定集合中某个索引的详细信息。\n",
    "\n",
    "    参数:\n",
    "    collection_name (str): 集合的名称。\n",
    "\n",
    "    返回:\n",
    "    None: 该函数仅打印索引信息，不返回任何值。\n",
    "    \"\"\"\n",
    "    # 查看集合的所有索引\n",
    "    indexes = milvus_client.list_indexes(\n",
    "        collection_name=collection_name  \n",
    "    )\n",
    "    print(f\"已经创建的索引：{indexes}\")\n",
    "    print()\n",
    "    # 查看索引信息\n",
    "    if indexes:\n",
    "        for index in indexes:\n",
    "            index_details = milvus_client.describe_index(\n",
    "                collection_name=collection_name,  \n",
    "                # 指定索引名称，这里假设使用第一个索引\n",
    "                index_name=index\n",
    "            )\n",
    "            print(f\"索引 {index} 详情：{index_details}\")\n",
    "            print()\n",
    "    else:\n",
    "        print(f\"集合 {collection_name} 中没有创建索引。\")\n",
    "\n",
    "# 示例\n",
    "show_index_info(collection_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 加载集合\n",
    "print(f\"正在加载集合：{collection_name}\")\n",
    "milvus_client.load_collection(collection_name=collection_name)\n",
    "\n",
    "# 验证加载状态\n",
    "print(milvus_client.get_load_state(collection_name=collection_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 定义稀疏向量搜索参数\n",
    "search_params_sparse_vectors = {\n",
    "    \"metric_type\": \"IP\",\n",
    "    \"params\": {\"drop_ratio_search\": 0.2},\n",
    "}\n",
    "\n",
    "# 执行向量搜索\n",
    "def vector_search(\n",
    "        query_vectors,\n",
    "        field_name,\n",
    "        search_params,\n",
    "        output_fields,\n",
    "    ):\n",
    "    # 向量搜索\n",
    "    res = milvus_client.search(\n",
    "        collection_name=collection_name,\n",
    "        # 指定查询向量。\n",
    "        data=query_vectors,\n",
    "        # 指定要搜索的向量字段\n",
    "        anns_field=field_name,\n",
    "        # 设置搜索参数\n",
    "        search_params=search_params,\n",
    "        output_fields=output_fields\n",
    "    )\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 打印向量搜索结果\n",
    "def print_vector_results(res):\n",
    "    for hits in res:\n",
    "        for hit in hits:\n",
    "            entity = hit.get(\"entity\")\n",
    "            print(f\"text: {entity['text']}\")\n",
    "            print(f\"distance: {hit['distance']:.3f}\")\n",
    "            print(\"-\"*50)\n",
    "        print(f\"数量：{len(hits)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 使用稀疏向量BM25搜索\n",
    "query1 = [\"人工智能如何影响汽车行业？\"]\n",
    "\n",
    "query_sparse_vectors_bm25 = bm25_ef.encode_queries(query1)\n",
    "\n",
    "field_name = \"sparse_vectors_bm25\"\n",
    "output_fields = [\"text\"]\n",
    "# 指定搜索的分区，或者过滤搜索\n",
    "res_sparse_vectors_bm25 = vector_search(query_sparse_vectors_bm25, field_name, search_params_sparse_vectors, output_fields)\n",
    "\n",
    "print_vector_results(res_sparse_vectors_bm25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 查看query1的分词结果\n",
    "print(analyzer(query1[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "分词结果只有“人工智能”一个词：\n",
    "\n",
    "['人工智能', '影响', '汽车行业']\n",
    "BM25的词汇表中虽然有“智能”这个词，但是并不包含“人工智能”、“影响”和“汽车行业”这些词，所以没有返回任何结果。\n",
    "\n",
    "我们把“人工智能”替换成“机器智能”，就可以搜索到了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 使用稀疏向量BM25搜索\n",
    "query2 = [\"机器智能如何影响汽车行业？\"]\n",
    "\n",
    "query_sparse_vectors_bm25 = bm25_ef.encode_queries(query2)\n",
    "\n",
    "field_name = \"sparse_vectors_bm25\"\n",
    "output_fields = [\"text\"]\n",
    "# 指定搜索的分区，或者过滤搜索\n",
    "res_sparse_vectors_bm25 = vector_search(query_sparse_vectors_bm25, field_name, search_params_sparse_vectors, output_fields)\n",
    "\n",
    "print_vector_results(res_sparse_vectors_bm25)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "而且，这次还搜索到了包含“机器学习”的句子。\n",
    "\n",
    "text: 机器智能的未来充满潜力。\n",
    "distance: 2.054\n",
    "--------------------------------------------------\n",
    "text: 大数据支持是机器智能发展的关键。\n",
    "distance: 1.752\n",
    "--------------------------------------------------\n",
    "text: 机器学习正在改变我们的生活方式。\n",
    "distance: 0.788\n",
    "--------------------------------------------------\n",
    "数量：3\n",
    "这是因为分词时把“机器智能“分成了“机器”和“智能”两个词，所以能搜索到更多句子。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 查看query2的分词结果\n",
    "print(analyzer(query2[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来，我们使用 splade 搜索，看看和 BM25的搜索结果有什么不同。\n",
    "\n",
    "先定义一个打印结果的函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "# 打印向量搜索结果\n",
    "def print_vector_results_en(res):\n",
    "    for hits in res:\n",
    "        for hit in hits:\n",
    "            entity = hit.get(\"entity\")\n",
    "            print(f\"text_en: {entity['text_en']}\")\n",
    "            print(f\"distance: {hit['distance']:.3f}\")\n",
    "            print(\"-\"*50)\n",
    "        print(f\"数量：{len(hits)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "ini"
    }
   },
   "outputs": [],
   "source": [
    "query1_en = [\"How does artificial intelligence affect the automotive industry?\"]\n",
    "\n",
    "query_sparse_vectors_splade = splade_ef.encode_queries(query1_en)\n",
    "\n",
    "field_name = \"sparse_vectors_splade\"\n",
    "output_fields = [\"text_en\"]\n",
    "res_sparse_vectors_splade = vector_search(query_sparse_vectors_splade, field_name, search_params_sparse_vectors, output_fields)\n",
    "\n",
    "print_vector_results_en(res_sparse_vectors_splade)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "比较 BM25 和 splade 的搜索结果，我们很容易发现它们之间的区别。splade 的文档集合中并不包含“artificial intelligence”这个词，但是由于它具有“举一反三”的能力，仍然搜索到了包含“AI”、“machine intelligence”以及“Autonomous”的句子，返回了更多结果（其实是返回了所有文档）。\n",
    ">```text\n",
    ">text_en: The future of machine intelligence is full of potential.\n",
    ">distance: 10.020\n",
    ">--------------------------------------------------\n",
    ">text_en: Big data support is key to the development of machine intelligence.\n",
    ">distance: 8.232\n",
    ">--------------------------------------------------\n",
    ">text_en: AI can help doctors diagnose diseases.\n",
    ">distance: 7.291\n",
    ">--------------------------------------------------\n",
    ">text_en: Autonomous driving relies on advanced algorithms.\n",
    ">distance: 7.213\n",
    ">--------------------------------------------------\n",
    ">text_en: Production efficiency can be improved through automation technology.\n",
    ">distance: 6.999\n",
    ">--------------------------------------------------\n",
    ">text_en: Machine learning is changing our way of life.\n",
    ">distance: 6.863\n",
    ">--------------------------------------------------\n",
    ">text_en: Data analysis technology is widely applied in the financial field.\n",
    ">distance: 5.064\n",
    ">--------------------------------------------------\n",
    ">text_en: The quantum tunneling effect allows electrons to pass through potential barriers that classical mechanics consider impassable, which has important applications in semiconductor devices.\n",
    ">distance: 3.695\n",
    ">--------------------------------------------------\n",
    ">text_en: Deep learning performs exceptionally well in image recognition.\n",
    ">distance: 3.464\n",
    ">--------------------------------------------------\n",
    ">text_en: Natural language processing is an important field in computer science.\n",
    ">distance: 3.044\n",
    ">--------------------------------------------------\n",
    "数量：10\n",
    "```\n",
    "\n",
    "如果把查询中的“artificial intelligence”替换成“machine intelligence”，仍然会返回所有结果，但是权重有所不同。\n",
    "\n",
    ">```text\n",
    ">text_en: The future of machine intelligence is full of potential.\n",
    ">distance: 15.128\n",
    ">--------------------------------------------------\n",
    ">text_en: Big data support is key to the development of machine intelligence.\n",
    ">distance: 12.945\n",
    ">--------------------------------------------------\n",
    ">text_en: Machine learning is changing our way of life.\n",
    ">distance: 12.763\n",
    ">--------------------------------------------------\n",
    ">text_en: Production efficiency can be improved through automation technology.\n",
    ">distance: 7.446\n",
    ">--------------------------------------------------\n",
    ">text_en: AI can help doctors diagnose diseases.\n",
    ">distance: 6.055\n",
    ">--------------------------------------------------\n",
    ">text_en: Autonomous driving relies on advanced algorithms.\n",
    ">distance: 5.309\n",
    ">--------------------------------------------------\n",
    ">text_en: Data analysis technology is widely applied in the financial field.\n",
    ">distance: 4.857\n",
    ">--------------------------------------------------\n",
    ">text_en: The quantum tunneling effect allows electrons to pass through potential barriers that classical mechanics consider impassable, which has important applications in semiconductor devices.\n",
    ">distance: 3.356\n",
    ">--------------------------------------------------\n",
    ">text_en: Deep learning performs exceptionally well in image recognition.\n",
    ">distance: 3.317\n",
    ">--------------------------------------------------\n",
    ">text_en: Natural language processing is an important field in computer science.\n",
    ">distance: 2.688\n",
    ">--------------------------------------------------\n",
    ">数量：10\n",
    ">```\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
